'''
model:
    pose env 
data:
    data_poseenv_control_100.json

额外功能
    在保存ckp时，使用json中val部分数据进行生成操作
'''

import os,sys
osa = os.path.abspath
osd = os.path.dirname
cur_dir = osd(osa(__file__))
par_dir = osd(cur_dir)
sys.path.insert(0,par_dir)

import warnings
warnings.filterwarnings('ignore')
import copy
import logging
import math
import os
import random
import shutil
from pathlib import Path

import accelerate
import numpy as np
import torch
import transformers
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import DistributedType, ProjectConfiguration, set_seed

from datasets import load_dataset
from packaging import version
from peft import LoraConfig, set_peft_model_state_dict
from peft.utils import get_peft_model_state_dict
from PIL import Image
from tqdm.auto import tqdm

import diffusers
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, \
     FluxTransformer2DModel,FluxKontextPipeline
from diffusers.optimization import get_scheduler
from diffusers.training_utils import (
    cast_training_params,
    compute_density_for_timestep_sampling,
    compute_loss_weighting_for_sd3,
    free_memory,
)
from diffusers.utils.torch_utils import is_compiled_module
from diffusers.image_processor import VaeImageProcessor
# from image_gen_aux import DepthPreprocessor


import pdb

from train_flux_parse import parse_args
from util_flux import process_img_1024
target_shape = (1024,1024)


# 额外把 图像处理的 clip 以及 lotus加载进来
# 使用another device 

logger = get_logger(__name__)

NORM_LAYER_PREFIXES = ["norm_q", "norm_k", "norm_added_q", "norm_added_k"]

input_prompt = "refcontrolPoseEnv change pose and env to photo with reference from right side," + \
            "change person to photo with refence from left side"

def encode_images(pixels: torch.Tensor, vae: torch.nn.Module, weight_dtype):
    pixel_latents = vae.encode(pixels.to(vae.dtype)).latent_dist.sample()
    pixel_latents = (pixel_latents - vae.config.shift_factor) * vae.config.scaling_factor
    return pixel_latents.to(weight_dtype)


def get_train_dataset(args, accelerator):
    dataset = None
    if args.dataset_name is not None:
        dataset = load_dataset(
            args.dataset_name,
            data_files=args.dataset_config_name,
        )
    if args.jsonl_for_train is not None:
        # load from json
        dataset = load_dataset("json", 
                                field='train',
                               data_files=args.jsonl_for_train, 
                               cache_dir=args.cache_dir)
        dataset = dataset.flatten_indices()
    # Preprocessing the datasets.
    # We need to tokenize inputs and targets.
    # column_names = dataset["train"].column_names
    '''
    will_save = {
        'train':[
                {
                    'prompt':input_prompt,
                    'input':input_img_save_path,
                }....
        ],
        'val':[....],
        }
    DatasetDict({
        train: Dataset({
            features: ['prompt', 'input'],
            num_rows: 100
        })
    })
    '''
    with accelerator.main_process_first():
        train_dataset = dataset['train'].shuffle(seed=args.seed)
        if args.max_train_samples is not None:
            train_dataset = train_dataset.select(range(args.max_train_samples))
    return train_dataset


def prepare_train_dataset(dataset, accelerator,image_processor, mask_processor):
    def preprocess_train(examples):
        
        concat_posenv = examples['input'][0]
        # 前 2048*1024 = target 
        
        concat_posenv_img = Image.open(concat_posenv)
        img_target_pil = concat_posenv_img.crop( (0,0,2048,1024) )
        w,h = concat_posenv_img.size
        img_ori_pil = concat_posenv_img.crop( (w-2048,h-1024 ,w,h ) )
        
        img_target = [
            img_target_pil
        ]
        img_ori = [
            img_ori_pil
        ]
        
        
        img_target_np = [image_processor.preprocess(image, 
                            height=args.resolution_height, 
                            width=args.resolution_width) 
                        for image in img_target]
        img_ori_np = [image_processor.preprocess(image, 
                            height=args.resolution_height, 
                            width=args.resolution_width) 
                        for image in img_ori]
        
        examples["img_target_pil"] = img_target   
        examples['img_ori_pil'] = img_ori
        examples["img_target"] = img_target_np   
        examples['img_ori'] = img_ori_np
        
        return examples

    with accelerator.main_process_first():
        dataset = dataset.with_transform(preprocess_train)



    return dataset


def collate_fn(examples):
    img_ori = torch.stack([example["img_ori"] for example in examples])
    img_ori = img_ori.to(memory_format=torch.contiguous_format).float()
    img_target = torch.stack([example["img_target"] for example in examples])
    img_target = img_target.to(memory_format=torch.contiguous_format).float()
    
    return {
            "captions": examples[0]['prompt'],
            
            "img_ori_pil": examples[0]['img_ori_pil'], 
            "img_target_pil": examples[0]['img_target_pil'], 
            
            "img_ori": img_ori, 
            "img_target": img_target,
            }


def main(args):
    
    if args.use_lora_bias and args.gaussian_init_lora:
        raise ValueError("`gaussian` LoRA init scheme isn't supported when `use_lora_bias` is True.")

    logging_out_dir = Path(args.output_dir, args.logging_dir)


    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, 
                                                      logging_dir=str(logging_out_dir))

    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
        log_with=args.report_to,
        project_config=accelerator_project_config,
        # debug=True
        # kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=True)]
    )

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        # DEBUG, INFO, WARNING, ERROR, CRITICAL
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)

    if accelerator.is_local_main_process:
        transformers.utils.logging.set_verbosity_warning()
        diffusers.utils.logging.set_verbosity_info()
    else:
        transformers.utils.logging.set_verbosity_error()
        diffusers.utils.logging.set_verbosity_error()

    # If passed along, set the training seed now.
    if args.seed is not None:
        set_seed(args.seed)

    # Handle the repository creation
    if accelerator.is_main_process:
        if args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)


    # cast down and move to the CPU
    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16
    
    text_encoding_pipeline = FluxKontextPipeline.from_pretrained(
        args.pretrained_model_name_or_path, transformer=None, vae=None, torch_dtype=weight_dtype
    ).to(accelerator.device)
    # Load models. We will load the text encoders later in a pipeline to compute
    # embeddings.
    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="vae",
        revision=args.revision,
        variant=args.variant,
    )
    vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
    image_processor = VaeImageProcessor( # vae channel 4 scale 16
        vae_scale_factor=vae_scale_factor * 2)
    latent_channels = vae.config.latent_channels
    mask_processor = VaeImageProcessor( # vae channel 16 scale 16
            vae_scale_factor=vae_scale_factor * 2,
            vae_latent_channels=latent_channels,
            do_normalize=False,
            do_binarize=True,
            do_convert_grayscale=True,
        )
    flux_transformer = FluxTransformer2DModel.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="transformer",
        revision=args.revision,
        variant=args.variant,
    )
    logger.info("All models loaded successfully")

    print("loaded pipeline")

    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="scheduler",
    )
    noise_scheduler_copy = copy.deepcopy(noise_scheduler)
    vae.requires_grad_(False)
    flux_transformer.requires_grad_(False)

    # let's not move the VAE to the GPU yet.
    vae.to(device=accelerator.device,dtype=torch.float32)  # keep the VAE in float32.

    # enable image inputs
    with torch.no_grad():
        initial_input_channels = flux_transformer.config.in_channels # kontext 64
        initial_output_channels = flux_transformer.config.out_channels
        new_linear = torch.nn.Linear( # ori_in 64 now_in 74 out 3072
            # flux_transformer.x_embedder.in_features * 3,
            flux_transformer.x_embedder.in_features, # 64
            flux_transformer.x_embedder.out_features,
            bias=flux_transformer.x_embedder.bias is not None,
            dtype=flux_transformer.dtype,
            device=flux_transformer.device,
        )


        new_linear.weight.zero_()
        new_linear.weight[:, :initial_input_channels].copy_(flux_transformer.x_embedder.weight)
        if flux_transformer.x_embedder.bias is not None:
            new_linear.bias.copy_(flux_transformer.x_embedder.bias)
        flux_transformer.x_embedder = new_linear

    assert torch.all(flux_transformer.x_embedder.weight[:, initial_input_channels:].data == 0)
    flux_transformer.register_to_config(
        # in_channels=initial_input_channels * 3, 
        in_channels=initial_input_channels, 
        out_channels=initial_output_channels)

    flux_transformer.to(accelerator.device, dtype=weight_dtype)

    if args.train_norm_layers:
        for name, param in flux_transformer.named_parameters():
            if any(k in name for k in NORM_LAYER_PREFIXES):
                param.requires_grad = True

    if args.lora_layers is not None:
        if args.lora_layers != "all-linear":
            target_modules = [layer.strip() for layer in args.lora_layers.split(",")]
            # add the input layer to the mix.
            if "x_embedder" not in target_modules:
                target_modules.append("x_embedder")
        elif args.lora_layers == "all-linear":
            target_modules = set()
            for name, module in flux_transformer.named_modules():
                if isinstance(module, torch.nn.Linear):
                    target_modules.add(name)
            target_modules = list(target_modules)
    else:
        target_modules = [
            "x_embedder",
            "attn.to_k",
            "attn.to_q",
            "attn.to_v",
            "attn.to_out.0",
            "attn.add_k_proj",
            "attn.add_q_proj",
            "attn.add_v_proj",
            "attn.to_add_out",
            "ff.net.0.proj",
            "ff.net.2",
            "ff_context.net.0.proj",
            "ff_context.net.2",
        ]
    transformer_lora_config = LoraConfig(
        r=args.rank,
        lora_alpha=args.rank,
        init_lora_weights="gaussian" if args.gaussian_init_lora else True,
        target_modules=target_modules,
        lora_bias=args.use_lora_bias,
    )
    flux_transformer.add_adapter(transformer_lora_config)


    def prepare_mask_latents(
        mask,
        masked_image,
        batch_size,
        num_channels_latents,
        num_images_per_prompt,
        height,
        width,
        dtype,
        device,
    ):
        # 1. calculate the height and width of the latents
        # VAE applies 8x compression on images but we must also account for packing which requires
        # latent height and width to be divisible by 2.
        height = 2 * (int(height) // (vae_scale_factor * 2))
        width = 2 * (int(width) // (vae_scale_factor * 2))

        # 2. encode the masked image
        if masked_image.shape[1] == num_channels_latents:
            masked_image_latents = masked_image
        else:
            masked_image_latents = vae.encode(masked_image.to(device=vae.device, dtype=vae.dtype)).latent_dist.sample()

        masked_image_latents = (masked_image_latents - vae.config.shift_factor) * vae.config.scaling_factor
        masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)

        # 3. duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method
        batch_size = batch_size * num_images_per_prompt
        if mask.shape[0] < batch_size:
            if not batch_size % mask.shape[0] == 0:
                raise ValueError(
                    "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
                    f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"
                    " of masks that you pass is divisible by the total requested batch size."
                )
            mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)
        if masked_image_latents.shape[0] < batch_size:
            if not batch_size % masked_image_latents.shape[0] == 0:
                raise ValueError(
                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"
                    f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."
                    " Make sure the number of images that you pass is divisible by the total requested batch size."
                )
            masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)

        # 5.resize mask to latents shape we we concatenate the mask to the latents
        mask = mask[:, 0, :, :]  # batch_size, 8 * height, 8 * width (mask has not been 8x compressed)
        mask = mask.view(
            batch_size, height, vae_scale_factor, width, vae_scale_factor
        )  # batch_size, height, 8, width, 8
        mask = mask.permute(0, 2, 4, 1, 3)  # batch_size, 8, 8, height, width
        mask = mask.reshape(
            batch_size, vae_scale_factor * vae_scale_factor, height, width
        )  # batch_size, 8*8, height, width
        mask = mask.to(device=device, dtype=dtype)

        return mask, masked_image_latents
    
    def unwrap_model(model):
        model = accelerator.unwrap_model(model)
        model = model._orig_mod if is_compiled_module(model) else model
        return model

    # `accelerate` 0.16.0 will have better support for customized saving
    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):

        def save_model_hook(models, weights, output_dir):
            if accelerator.is_main_process:
                transformer_lora_layers_to_save = None

                embder_count = 1
                for model in models:
                    if isinstance(unwrap_model(model), type(unwrap_model(flux_transformer))):
                        model = unwrap_model(model)
                        transformer_lora_layers_to_save = get_peft_model_state_dict(model)
                        if args.train_norm_layers:
                            transformer_norm_layers_to_save = {
                                f"transformer.{name}": param
                                for name, param in model.named_parameters()
                                if any(k in name for k in NORM_LAYER_PREFIXES)
                            }
                            transformer_lora_layers_to_save = {
                                **transformer_lora_layers_to_save,
                                **transformer_norm_layers_to_save,
                            }
                    # elif isinstance(unwrap_model(model), ReduxImageEncoder):
                    #     torch.save(unwrap_model(model).state_dict(),
                    #           os.path.join(output_dir, f"image_encoder_{embder_count}.bin"))
                    #     embder_count+=1
                    else:
                        raise ValueError(f"unexpected save model: {model.__class__}")

                    # make sure to pop weight so that corresponding model is not saved again
                    if weights:
                        weights.pop()

                FluxKontextPipeline.save_lora_weights(
                    output_dir,
                    transformer_lora_layers=transformer_lora_layers_to_save,
                )

        def load_model_hook(models, input_dir):
            transformer_ = None

            if not accelerator.distributed_type == DistributedType.DEEPSPEED:
                while len(models) > 0:
                    model = models.pop()

                    if isinstance(model, type(unwrap_model(flux_transformer))):
                        transformer_ = model
                    else:
                        raise ValueError(f"unexpected save model: {model.__class__}")
            else:
                transformer_ = FluxTransformer2DModel.from_pretrained(
                    args.pretrained_model_name_or_path, subfolder="transformer"
                ).to(accelerator.device, weight_dtype)

                # Handle input dimension doubling before adding adapter
                with torch.no_grad():
                    initial_input_channels = transformer_.config.in_channels
                    new_linear = torch.nn.Linear(
                        transformer_.x_embedder.in_features,
                        transformer_.x_embedder.out_features,
                        bias=transformer_.x_embedder.bias is not None,
                        dtype=transformer_.dtype,
                        device=transformer_.device,
                    )
                    new_linear.weight.zero_()
                    new_linear.weight[:, :initial_input_channels].copy_(transformer_.x_embedder.weight)
                    if transformer_.x_embedder.bias is not None:
                        new_linear.bias.copy_(transformer_.x_embedder.bias)
                    transformer_.x_embedder = new_linear
                    transformer_.register_to_config(in_channels=initial_input_channels * 2)

                transformer_.add_adapter(transformer_lora_config)

            lora_state_dict = FluxKontextPipeline.lora_state_dict(input_dir)
            transformer_lora_state_dict = {
                f'{k.replace("transformer.", "")}': v
                for k, v in lora_state_dict.items()
                if k.startswith("transformer.") and "lora" in k
            }
            incompatible_keys = set_peft_model_state_dict(
                transformer_, transformer_lora_state_dict, adapter_name="default"
            )
            if incompatible_keys is not None:
                # check only for unexpected keys
                unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
                if unexpected_keys:
                    logger.warning(
                        f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
                        f" {unexpected_keys}. "
                    )
            if args.train_norm_layers:
                transformer_norm_state_dict = {
                    k: v
                    for k, v in lora_state_dict.items()
                    if k.startswith("transformer.") and any(norm_k in k for norm_k in NORM_LAYER_PREFIXES)
                }
                transformer_._transformer_norm_layers = FluxKontextPipeline._load_norm_into_transformer(
                    transformer_norm_state_dict,
                    transformer=transformer_,
                    discard_original_layers=False,
                )

            # Make sure the trainable params are in float32. This is again needed since the base models
            # are in `weight_dtype`. More details:
            # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
            if args.mixed_precision == "fp16":
                models = [transformer_]
                # only upcast trainable parameters (LoRA) into fp32
                cast_training_params(models)

        accelerator.register_save_state_pre_hook(save_model_hook)
        accelerator.register_load_state_pre_hook(load_model_hook)

    # Make sure the trainable params are in float32.
    if args.mixed_precision == "fp16":
        models = [flux_transformer]
        # only upcast trainable parameters (LoRA) into fp32
        cast_training_params(models, dtype=torch.float32)

    if args.gradient_checkpointing:
        flux_transformer.enable_gradient_checkpointing()

    # Enable TF32 for faster training on Ampere GPUs,
    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
    if args.allow_tf32:
        torch.backends.cuda.matmul.allow_tf32 = True

    if args.scale_lr:
        args.learning_rate = (
            args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
        )

    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
    optimizer_class = torch.optim.AdamW


    # Optimization parameters
    transformer_lora_parameters = list(filter(lambda p: p.requires_grad, flux_transformer.parameters()))
                                        # list( dpth_embder.parameters() ) 
    # print(transformer_lora_parameters)
    optimizer = optimizer_class(
        transformer_lora_parameters,
        lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

    # Prepare dataset and dataloader.
    train_dataset = get_train_dataset(args, accelerator)
    train_dataset = prepare_train_dataset(train_dataset, accelerator, image_processor, mask_processor)
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        shuffle=True,
        collate_fn=collate_fn,
        batch_size=args.train_batch_size,
        num_workers=args.dataloader_num_workers,
    )

    # Scheduler and math around the number of training steps.
    # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation.
    if args.max_train_steps is None:
        len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)
        num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)
        num_training_steps_for_scheduler = (
            args.num_train_epochs * num_update_steps_per_epoch * accelerator.num_processes
        )
    else:
        num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes

    lr_scheduler = get_scheduler(
        args.lr_scheduler,
        optimizer=optimizer,
        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
        num_training_steps=args.max_train_steps * accelerator.num_processes,
        num_cycles=args.lr_num_cycles,
        power=args.lr_power,
    )
    # Prepare everything with our `accelerator`.
    flux_transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        flux_transformer, optimizer, train_dataloader, lr_scheduler  )

    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        if num_training_steps_for_scheduler != args.max_train_steps * accelerator.num_processes:
            logger.warning(
                f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match "
                f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. "
                f"This inconsistency may result in the learning rate scheduler not functioning properly."
            )
    # Afterwards we recalculate our number of training epochs
    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

    # We need to initialize the trackers we use, and also store our configuration.
    # The trackers initializes automatically on the main process.
    if accelerator.is_main_process:
        tracker_config = dict(vars(args))
        accelerator.init_trackers(args.tracker_project_name, config=tracker_config)

    # Train!
    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")
    global_step = 0
    first_epoch = 0

    # Potentially load in the weights and states from a previous save
    if args.resume_from_checkpoint:
        if args.resume_from_checkpoint != "latest":
            path = os.path.basename(args.resume_from_checkpoint)
        else:
            # Get the most recent checkpoint
            dirs = os.listdir(args.output_dir)
            dirs = [d for d in dirs if d.startswith("checkpoint")]
            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
            path = dirs[-1] if len(dirs) > 0 else None

        if path is None:
            logger.info(f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.")
            args.resume_from_checkpoint = None
            initial_global_step = 0
        else:
            logger.info(f"Resuming from checkpoint {path}")
            accelerator.load_state(os.path.join(args.output_dir, path))
            global_step = int(path.split("-")[1])

            initial_global_step = global_step
            first_epoch = global_step // num_update_steps_per_epoch
    else:
        initial_global_step = 0

    progress_bar = tqdm(
        range(0, args.max_train_steps),
        initial=initial_global_step,
        desc="Steps",
        # Only show the progress bar once on each machine.
        disable=not accelerator.is_local_main_process,
    )

    def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
        sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
        schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device)
        timesteps = timesteps.to(accelerator.device)
        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]

        sigma = sigmas[step_indices].flatten()
        while len(sigma.shape) < n_dim:
            sigma = sigma.unsqueeze(-1)
        return sigma

    from utils.util_flux import horizontal_concat_images

    prompt_embeds, pooled_prompt_embeds, text_ids = None,None,None
    

    for epoch in range(first_epoch, args.num_train_epochs):
        flux_transformer.train()
        for step, batch in enumerate(train_dataloader):
            torch.cuda.empty_cache()
            with accelerator.accumulate(flux_transformer):
                

                # test_ori = batch['img_ori_pil']
                # test_target = batch['img_target_pil']
                # horizontal_concat_images( [test_ori,test_target],
                #                          gap=0,line_width=0, ).save('tmp_posenv.jpg')
                
                # text encoding.
                if prompt_embeds is None:
                    captions = batch["captions"]
                    with torch.no_grad():
                        prompt_embeds, pooled_prompt_embeds, text_ids = text_encoding_pipeline.encode_prompt(
                            captions, prompt_2=captions,
                        )


                img_origin_pixel_latents= encode_images( batch['img_ori'].squeeze(1).to(  #torch.Size([1, 16, 128, 256])
                                                        device=accelerator.device, 
                                                        dtype=weight_dtype) , vae , weight_dtype )


                # vae encode 1 16 128 128
                img_target_pixel_latents = encode_images(batch["img_target"].squeeze(1), 
                                                        vae, weight_dtype)
                

                # pdb.set_trace()
   

                # Sample a random timestep for each image
                # for weighting schemes where we sample timesteps non-uniformly
                bsz = img_origin_pixel_latents.shape[0]
                # num_images_per_prompt = 1
                # torch.Size([1, 64, 128, 256]) masked_image_latents
                
                # noise : torch.Size([1, 16, 128, 384]) 128*384/4=12288
                noise = torch.randn_like(img_target_pixel_latents, device=accelerator.device, dtype=weight_dtype)
                u = compute_density_for_timestep_sampling(
                    weighting_scheme=args.weighting_scheme,
                    batch_size=bsz,
                    logit_mean=args.logit_mean,
                    logit_std=args.logit_std,
                    mode_scale=args.mode_scale,
                )
                indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
                timesteps = noise_scheduler_copy.timesteps[indices].to(device=img_target_pixel_latents.device)

                # Add noise according to flow matching.
                sigmas = get_sigmas(timesteps, n_dim=img_target_pixel_latents.ndim, dtype=img_target_pixel_latents.dtype)
                noisy_model_input = (1.0 - sigmas) * img_target_pixel_latents + sigmas * noise

                # pack the latents.
                packed_noisy_model_input = FluxKontextPipeline._pack_latents(
                    noisy_model_input, # torch.Size([1, 12288, 64])
                    batch_size=bsz,
                    num_channels_latents=noisy_model_input.shape[1],
                    height=noisy_model_input.shape[2],
                    width=noisy_model_input.shape[3],
                )
                packed_ori_latents = FluxKontextPipeline._pack_latents(
                    img_origin_pixel_latents, # torch.Size([1, 4096, 64]) 
                    batch_size=bsz,
                    num_channels_latents=img_origin_pixel_latents.shape[1],
                    height=img_origin_pixel_latents.shape[2],
                    width=img_origin_pixel_latents.shape[3],
                )
                '''
                torch.Size([1, 16, 128, 128]) torch.Size([1, 4096, 64]) 16*4=64
                torch.Size([1, 16, 128, 128]) torch.Size([1, 4096, 64])
                
                default-FLUX-Fill : torch.Size([1, 12288, 384]) 12288/64=192
                default-FLUX-Depth : 128 = 64 ori + 64 control

                now = 64
                '''
                packed_noisy_model_final_input = torch.cat((packed_noisy_model_input, 
                                                            packed_ori_latents,
                                                        ), dim=-2) # 在 h*w 上拼接
                

                # latent image ids for RoPE. torch.Size([h*w, 3])
                image_ids = FluxKontextPipeline._prepare_latent_image_ids( # torch.Size([8192, 3])
                    bsz,
                    img_origin_pixel_latents.shape[2] // 2,
                    img_origin_pixel_latents.shape[3] // 2,
                    accelerator.device,
                    weight_dtype,
                )
                image_ids[..., 0] = 1
                latent_image_ids = FluxKontextPipeline._prepare_latent_image_ids( # torch.Size([8192, 3])
                    bsz,
                    noisy_model_input.shape[2] // 2,
                    noisy_model_input.shape[3] // 2,
                    accelerator.device,
                    weight_dtype,
                )
                latent_image_ids = torch.cat([latent_image_ids, image_ids], dim=0)

                # handle guidance
                if unwrap_model(flux_transformer).config.guidance_embeds:
                    guidance_vec = torch.full(
                        (bsz,),
                        args.guidance_scale,
                        device=noisy_model_input.device,
                        dtype=weight_dtype,
                    )
                else:
                    guidance_vec = None

                # Predict.
                model_pred = flux_transformer( # model_pred torch.Size([1, 4096*2, 64])
                    hidden_states=packed_noisy_model_final_input,
                    timestep=timesteps / 1000,
                    guidance=guidance_vec,
                    pooled_projections=pooled_prompt_embeds, # 1 768
                    encoder_hidden_states=prompt_embeds,  # 1 1241 4096
                    txt_ids=text_ids, # torch.Size([1241, 3]) txt_id shape[0] = prompt shape[1]
                    # prompt = prompt,
                    img_ids=latent_image_ids, # torch.Size([4096*2, 3])
                    return_dict=False,
                )[0]
                model_pred = model_pred[:, : packed_noisy_model_input.size(1)] # torch.Size([1, 16384, 64])
                # pdb.set_trace()
                model_pred = FluxKontextPipeline._unpack_latents(
                    model_pred,
                    height=noisy_model_input.shape[2] * vae_scale_factor,
                    width=noisy_model_input.shape[3] * vae_scale_factor,
                    vae_scale_factor=vae_scale_factor,
                )
                # these weighting schemes use a uniform timestep sampling
                # and instead post-weight the loss
                weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, 
                                                           sigmas=sigmas)

                # flow-matching loss
                target = noise - img_target_pixel_latents # torch.Size([1, 16, 128, 256])
                # target = pixel_latents - noise
                loss = torch.mean(
                    (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
                    1,
                )
                loss = loss.mean()
                accelerator.backward(loss)

                if accelerator.sync_gradients:
                    params_to_clip = flux_transformer.parameters()
                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                progress_bar.update(1)
                global_step += 1

                # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues.
                if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
                    # if True or global_step % args.checkpointing_steps == 0:
                    if global_step % args.checkpointing_steps == 0:

                        save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                        accelerator.save_state(save_path)
                        logger.info(f"Saved state to {save_path}")

                        # 保存时 进行 额外的 val 操作
                        # val 的 也 存入 output_dir
                        val_save_dir = os.path.join(args.output_dir, f"val_{global_step}")

                        def val_transformer(flux_transformer, 
                                        args, accelerator, val_save_dir,
                                        prompt_embeds, pooled_prompt_embeds ):
                            os.makedirs(val_save_dir,exist_ok=True)

                            device = 'cuda:1'

                            weight_dtype = flux_transformer.dtype

                            flux_transformer = accelerator.unwrap_model(flux_transformer)
                            pipeline = FluxKontextPipeline.from_pretrained(
                                args.pretrained_model_name_or_path,
                                transformer=flux_transformer,
                                torch_dtype=weight_dtype,
                            )
                            pipeline.to(device)
                            pipeline.set_progress_bar_config(disable=True)

                            generator = torch.Generator(device=device)

                            autocast_ctx = torch.autocast(device, weight_dtype)
                            # 从 json 中直接读取 val img path
                            json_file = args.jsonl_for_train
                            import json
                            with open(json_file,encoding='utf-8') as f:
                                val_data = json.load(f)['val'] # list
                            for dict_data in val_data:
                                prom,img_path = dict_data.values()

                                concat_posenv_img = Image.open(img_path)
                                # img_target_pil = concat_posenv_img.crop( (0,0,2048,1024) )
                                w,h = concat_posenv_img.size
                                img_ori_pil = concat_posenv_img.crop( (w-2048,h-1024 ,w,h ) ) # 输入

                                with autocast_ctx:
                                    image = pipeline(
                                        # prompt=validation_prompt,
                                        image=img_ori_pil,
                                        num_inference_steps=50,
                                        guidance_scale=args.guidance_scale,
                                        generator=generator,
                                        max_sequence_length=512,
                                        height=img_ori_pil.height,
                                        width=img_ori_pil.width,
                                        prompt_embeds=prompt_embeds,
                                        pooled_prompt_embeds=pooled_prompt_embeds,
                                    ).images[0]
                                # image = image.resize((args.resolution, args.resolution))
                                pdb.set_trace()
                                image.save( os.path.join( val_save_dir, os.path.basename(img_path) ) )

                                break
                            print('val end..........')
                                
                        # val_transformer(flux_transformer, args, accelerator,
                        #                 val_save_dir,
                        #                 prompt_embeds, pooled_prompt_embeds)
            logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
            progress_bar.set_postfix(**logs)
            accelerator.log(logs, step=global_step)

            if global_step >= args.max_train_steps:
                break

    # Create the pipeline using using the trained modules and save it.
    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        flux_transformer = unwrap_model(flux_transformer)
        if args.upcast_before_saving:
            flux_transformer.to(torch.float32)
        transformer_lora_layers = get_peft_model_state_dict(flux_transformer)
        if args.train_norm_layers:
            transformer_norm_layers = {
                f"transformer.{name}": param
                for name, param in flux_transformer.named_parameters()
                if any(k in name for k in NORM_LAYER_PREFIXES)
            }
            transformer_lora_layers = {**transformer_lora_layers, **transformer_norm_layers}
        FluxKontextPipeline.save_lora_weights(
            save_directory=args.output_dir,
            transformer_lora_layers=transformer_lora_layers,
        )

        del flux_transformer
        del vae
        torch.cuda.empty_cache()
        free_memory()

    accelerator.end_training()


if __name__ == "__main__":
    args = parse_args()
    main(args)
